import tensorflow as tf
import config as cfg


# for ty in ["duration", "acoustic"]:
#     for phase in ["train", "test"]

class DataSet(object):

    def __init__(self, l_d, batch_size):
        x_d_train = tf.data.Dataset.from_generator(
            lambda: l_d.x['duration']['train'], tf.float32)
        y_d_train = tf.data.Dataset.from_generator(
            lambda: l_d.y['duration']['train'], tf.float32)
        ul_d_train = tf.data.Dataset.from_generator(
            lambda: l_d.utt_lengths['duration']['train'], tf.int32)
        x_a_train = tf.data.Dataset.from_generator(
            lambda: l_d.x['acoustic']['train'], tf.float32)
        y_a_train = tf.data.Dataset.from_generator(
            lambda: l_d.y['acoustic']['train'], tf.float32)
        ul_a_train = tf.data.Dataset.from_generator(
            lambda: l_d.utt_lengths['acoustic']['train'], tf.int32)
        self.sim_train_set = tf.data.Dataset.zip(
            (x_d_train, y_d_train, ul_d_train,
             x_a_train, y_a_train, ul_a_train))
        self.sim_train_set = self.sim_train_set.padded_batch(batch_size,
                                                             padded_shapes=(
                                                                 [None, None],
                                                                 [None, None],
                                                                 [],
                                                                 [None,
                                                                  cfg.a_in],
                                                                 [None, None],
                                                                 []))
        self.train_set = tf.data.Dataset.zip((x_d_train, y_d_train, ul_d_train,
                                              x_a_train, y_a_train, ul_a_train))
        self.train_set = self.train_set.repeat(count=cfg.train_set_repeat_count)
        if cfg.is_shuffle_train:
            self.train_set = self.train_set.shuffle(1000,
                                                    reshuffle_each_iteration=True)
        self.train_set = self.train_set.padded_batch(batch_size, padded_shapes=(
            [None, None], [None, None], [], [None, cfg.a_in],
            [None, None], []))

        x_d_test = tf.data.Dataset.from_generator(
            lambda: l_d.x['duration']['test'], tf.float32)
        y_d_test = tf.data.Dataset.from_generator(
            lambda: l_d.y['duration']['test'], tf.float32)
        ul_d_test = tf.data.Dataset.from_generator(
            lambda: l_d.utt_lengths['duration']['test'], tf.int32)
        x_a_test = tf.data.Dataset.from_generator(
            lambda: l_d.x['acoustic']['test'], tf.float32)
        y_a_test = tf.data.Dataset.from_generator(
            lambda: l_d.y['acoustic']['test'], tf.float32)
        ul_a_test = tf.data.Dataset.from_generator(
            lambda: l_d.utt_lengths['acoustic']['test'], tf.int32)
        self.test_set = tf.data.Dataset.zip((x_d_test, y_d_test, ul_d_test,
                                             x_a_test, y_a_test, ul_a_test))
        self.test_set = self.test_set.padded_batch(batch_size, padded_shapes=(
            [None, None], [None, None], [], [None, cfg.a_in],
            [None, None], []))

        self.train_iterator = self.train_set.make_initializable_iterator()
        self.sim_train_iterator = self.sim_train_set.make_initializable_iterator()
        self.test_iterator = self.test_set.make_initializable_iterator()

    def train_gen(self):
        return self.train_iterator

    def sim_train_gen(self):
        return self.sim_train_iterator

    def test_gen(self):
        return self.test_iterator
